{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Group Chat with Customized Speaker Selection Method\n",
    "\n",
    "AG2 offers conversable agents powered by LLM, tool or human, which can be used to perform tasks collectively via automated chat. This framework allows tool use and human participation through multi-agent conversation.\n",
    "Please find documentation about this feature [here](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/conversable-agent/).\n",
    "\n",
    "In this notebook, we demonstrate how to pass a customized agent selection method to GroupChat. The customized function looks like this:\n",
    "\n",
    "```python\n",
    "def custom_speaker_selection_func(last_speaker, groupchat):\n",
    "    \"\"\"Define a customized speaker selection function.\n",
    "    A recommended way is to define a transition for each speaker in the groupchat.\n",
    "\n",
    "    Parameters:\n",
    "        - last_speaker: Agent\n",
    "            The last speaker in the group chat.\n",
    "        - groupchat: GroupChat\n",
    "            The GroupChat object\n",
    "    Return:\n",
    "        Return one of the following:\n",
    "        1. an `Agent` class, it must be one of the agents in the group chat.\n",
    "        2. a string from ['auto', 'manual', 'random', 'round_robin'] to select a default method to use.\n",
    "        3. None, which indicates the chat should be terminated.\n",
    "    \"\"\"\n",
    "    pass\n",
    "\n",
    "groupchat = autogen.GroupChat(\n",
    "    speaker_selection_method=custom_speaker_selection_func,\n",
    "    ...,\n",
    ")\n",
    "```\n",
    "The last speaker and the groupchat object are passed to the function. Commonly used variables from groupchat are `groupchat.messages` an `groupchat.agents`, which is the message history and the agents in the group chat respectively. You can access other attributes of the groupchat, such as `groupchat.allowed_speaker_transitions_dict` for pre-defined allowed_speaker_transitions_dict. \n",
    "\n",
    "````{=mdx}\n",
    ":::info Requirements\n",
    "Install `ag2`:\n",
    "```bash\n",
    "pip install ag2[openai]\n",
    "```\n",
    "\n",
    "For more information, please refer to the [installation guide](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/installing-ag2).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set your API Endpoint\n",
    "\n",
    "The [`LLMConfig.from_json`](https://docs.ag2.ai/latest/docs/api-reference/autogen/llm_config/LLMConfig/#autogen.llm_config.LLMConfig.from_json) function loads a list of configurations from an environment variable or a json file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import autogen\n",
    "\n",
    "filter_dict = {\"model\": [\"gpt-4\", \"gpt-4-1106-preview\"]}\n",
    "llm_config = autogen.LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", cache_seed=42, temperature=0, timeout=120).where(\n",
    "    **filter_dict\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/llm-configuration).\n",
    ":::\n",
    "````\n",
    "\n",
    "## Construct Agents\n",
    "\n",
    "- Planner: Give a plan and revise.\n",
    "- Admin: Human in the loop to approve or terminate the process.\n",
    "- Engineer: Retrieve papers from the internet by writing code.\n",
    "- Executor: Execute the code.\n",
    "- Scientist: Read the papers and write a summary.\n",
    "\n",
    "The pipeline is the following:\n",
    "\n",
    "1. The planner interact with Admin (user) to revise a plan. Only when the Admin types \"Approve\", we can move to the next step.\n",
    "2. The engineer will write code to retrieve papers from the internet. The code will be executed by executor.\n",
    "3. When the code is executed successfully, the scientist will read the papers and write a summary.\n",
    "4. The summary will be reviewed by the Admin and give comments. When the Admin types \"TERMINATE\", the process will be terminated.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "planner = autogen.AssistantAgent(\n",
    "    name=\"Planner\",\n",
    "    system_message=\"\"\"Planner. Suggest a plan. Revise the plan based on feedback from admin and critic, until admin approval.\n",
    "The plan may involve an engineer who can write code and a scientist who doesn't write code.\n",
    "Explain the plan first. Be clear which step is performed by an engineer, and which step is performed by a scientist.\n",
    "\"\"\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "engineer = autogen.AssistantAgent(\n",
    "    name=\"Engineer\",\n",
    "    system_message=\"\"\"Engineer. You follow an approved plan. You write python/shell code to solve tasks. Wrap the code in a code block that specifies the script type. The user can't modify your code. So do not suggest incomplete code which requires others to modify. Don't use a code block if it's not intended to be executed by the executor.\n",
    "Don't include multiple code blocks in one response. Do not ask others to copy and paste the result. Check the execution result returned by the executor.\n",
    "If the result indicates there is an error, fix the error and output the code again. Suggest the full code instead of partial code or code changes. If the error can't be fixed or if the task is not solved even after the code is executed successfully, analyze the problem, revisit your assumption, collect additional info you need, and think of a different approach to try.\n",
    "\"\"\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "scientist = autogen.AssistantAgent(\n",
    "    name=\"Scientist\",\n",
    "    system_message=\"\"\"Scientist. You follow an approved plan. You are able to categorize papers after seeing their abstracts printed. You don't write code.\"\"\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "user_proxy = autogen.UserProxyAgent(\n",
    "    name=\"Admin\",\n",
    "    system_message=\"A human admin. Interact with the planner to discuss the plan. Plan execution needs to be approved by this admin.\",\n",
    "    code_execution_config=False,\n",
    ")\n",
    "\n",
    "executor = autogen.UserProxyAgent(\n",
    "    name=\"Executor\",\n",
    "    system_message=\"Executor. Execute the code written by the engineer and report the result.\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    code_execution_config={\n",
    "        \"last_n_messages\": 3,\n",
    "        \"work_dir\": \"paper\",\n",
    "        \"use_docker\": False,\n",
    "    },  # Please set use_docker=True if docker is available to run the generated code. Using docker is safer than running the generated code directly.\n",
    ")\n",
    "\n",
    "\n",
    "from autogen import Agent\n",
    "\n",
    "\n",
    "def custom_speaker_selection_func(last_speaker: Agent, groupchat: autogen.GroupChat):\n",
    "    \"\"\"Define a customized speaker selection function.\n",
    "    A recommended way is to define a transition for each speaker in the groupchat.\n",
    "\n",
    "    Returns:\n",
    "        Return an `Agent` class or a string from ['auto', 'manual', 'random', 'round_robin'] to select a default method to use.\n",
    "    \"\"\"\n",
    "    messages = groupchat.messages\n",
    "\n",
    "    if len(messages) <= 1:\n",
    "        return planner\n",
    "\n",
    "    if last_speaker is user_proxy:\n",
    "        if \"Approve\" in messages[-1][\"content\"]:\n",
    "            # If the last message is approved, let the engineer to speak\n",
    "            return engineer\n",
    "        elif messages[-2][\"name\"] == \"Planner\":\n",
    "            # If it is the planning stage, let the planner to continue\n",
    "            return planner\n",
    "        elif messages[-2][\"name\"] == \"Scientist\":\n",
    "            # If the last message is from the scientist, let the scientist to continue\n",
    "            return scientist\n",
    "\n",
    "    elif last_speaker is planner:\n",
    "        # Always let the user to speak after the planner\n",
    "        return user_proxy\n",
    "\n",
    "    elif last_speaker is engineer:\n",
    "        if \"```python\" in messages[-1][\"content\"]:\n",
    "            # If the last message is a python code block, let the executor to speak\n",
    "            return executor\n",
    "        else:\n",
    "            # Otherwise, let the engineer to continue\n",
    "            return engineer\n",
    "\n",
    "    elif last_speaker is executor:\n",
    "        if \"exitcode: 1\" in messages[-1][\"content\"]:\n",
    "            # If the last message indicates an error, let the engineer to improve the code\n",
    "            return engineer\n",
    "        else:\n",
    "            # Otherwise, let the scientist to speak\n",
    "            return scientist\n",
    "\n",
    "    elif last_speaker is scientist:\n",
    "        # Always let the user to speak after the scientist\n",
    "        return user_proxy\n",
    "\n",
    "    else:\n",
    "        return \"random\"\n",
    "\n",
    "\n",
    "groupchat = autogen.GroupChat(\n",
    "    agents=[user_proxy, engineer, scientist, planner, executor],\n",
    "    messages=[],\n",
    "    max_round=20,\n",
    "    speaker_selection_method=custom_speaker_selection_func,\n",
    ")\n",
    "\n",
    "manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start Chat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_proxy.initiate_chat(\n",
    "    manager, message=\"Find a latest paper about gpt-4 on arxiv and find its potential applications in software.\"\n",
    ")\n",
    "# type exit to terminate the chat"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Introduce group chat with customized speaker selection method.",
   "tags": [
    "orchestration",
    "group chat"
   ]
  },
  "kernelspec": {
   "display_name": "flaml",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
